import os
import argparse


"""
encode raw yuv by libx265 and extract images from the encoded video
"""



parser = argparse.ArgumentParser()
parser.add_argument("-f", "--seq_folder", type = str, required = True)
parser.add_argument("-sw", "--seq_width", type = str, required = True)
parser.add_argument("-sh", "--seq_height", type = str, required = True)

parser.add_argument("-o", "--output_folder", type = str, default = "./")

parser.add_argument("-c", "--codec", type = str, default = "libx265")
parser.add_argument("--crf", type = str, default = "20")
parser.add_argument("--gop", type = str, default = "12")
parser.add_argument("--preset", type = str, default = "veryfast")
parser.add_argument("--tune", type = str, default = "zerolatency")

parser.add_argument("--rename", action="store_true")


args = parser.parse_args()

# args
seq_folder = args.seq_folder
sw, sh = args.seq_width, args.seq_height

output_folder = args.output_folder

codec, crf, gop = args.codec, args.crf, args.gop
preset, tune = args.preset, args.tune

rename = args.rename

# folders

def run_cmd(options):
    cmd = ""
    for opt in options:
        opt = opt.replace('\\','/')
        cmd += (opt + " ")
    
    print(cmd)
    os.system(cmd)


def extract_image_from_yuv(seq):
    input_video = os.path.join(seq_folder, seq)
    folder = os.path.join(output_folder, seq[:-4], "uncompressed")

    options = [
        "ffmpeg",
        "-pix_fmt", "yuv420p",
        "-s:v", f"{sw}x{sh}",
        "-i", input_video,
        "-f", "image2",
        f"{folder}/im%04d.png"
    ]

    run_cmd(options)


def encode_yuv_and_extract(seq):
    input_video = os.path.join(seq_folder, seq)
    intermediate_folder = os.path.join(output_folder, seq[:-4], "out")
    
    options = [
        "ffmpeg",

        "-pix_fmt", "yuv420p",
        "-s:v", f"{sw}x{sh}",
        "-i", input_video,

        "-c:v", codec,
        "-preset", preset,
        "-tune", tune,
        "-x265-params", f"crf={crf}:keyint={gop}:verbose=1:csv-log-level=1:csv={intermediate_folder}/report.csv",

        f"{intermediate_folder}/out.mkv"
    ]

    run_cmd(options)

    folder = os.path.join(output_folder, seq[:-4], "compressed")
    encoded_video = os.path.join(intermediate_folder, "out.mkv")


    options = [
        "ffmpeg",
        "-i", encoded_video,
        "-f", "image2",
        f"{folder}/im%04d.png"
    ]
    run_cmd(options)


def rename_compreesed_images(seq):
    import pandas as pd

    report_path = os.path.join(output_folder, seq[:-4], "out", "report.csv")
    compressed_folder = os.path.join(output_folder, seq[:-4], "compressed")
    uncompressed_folder = os.path.join(output_folder, seq[:-4], "uncompressed")


    df = pd.read_csv(report_path)
    for img in os.listdir(compressed_folder):
        id = int(img[2:6])
        bits = df.loc[id - 1, " Bits"]
        type = df.loc[id - 1, " Type"].strip()[0]

        old_name = os.path.join(compressed_folder, img)
        new_name = os.path.join(compressed_folder, img[:-4] + f"_{bits}_{type}.png")
        os.rename(old_name, new_name)


if not os.path.exists(output_folder):
    os.makedirs(output_folder)

for seq in os.listdir(seq_folder):
    if seq.endswith(".yuv"):
        folder = os.path.join(output_folder, seq[:-4])
        if not os.path.exists(folder):
            os.makedirs(os.path.join(folder, "uncompressed"))
            os.makedirs(os.path.join(folder, "out"))
            os.makedirs(os.path.join(folder, "compressed"))

            extract_image_from_yuv(seq)

            encode_yuv_and_extract(seq)

        if rename:
            rename_compreesed_images(seq)


